﻿/*
 *  Injectors - Address Translation Management
 *
 *  Copyright (C) 2014 LINK/2012 <dma_2012@hotmail.com>
 *
 *  This software is provided 'as-is', without any express or implied
 *  warranty. In no event will the authors be held liable for any damages
 *  arising from the use of this software.
 *
 *  Permission is granted to anyone to use this software for any purpose,
 *  including commercial applications, and to alter it and redistribute it
 *  freely, subject to the following restrictions:
 *
 *     1. The origin of this software must not be misrepresented; you must not
 *     claim that you wrote the original software. If you use this software
 *     in a product, an acknowledgment in the product documentation would be
 *     appreciated but is not required.
 *
 *     2. Altered source versions must be plainly marked as such, and must not be
 *     misrepresented as being the original software.
 *
 *     3. This notice may not be removed or altered from any source
 *     distribution.
 *
 */
#pragma once

#if !defined(INJECTOR_GVM_HAS_TRANSLATOR)
#error Missing INJECTOR_GVM_HAS_TRANSLATOR on compiler definitions
#endif

 /*
  *  This is a quick solution for address translations if you're too lazy to implement a proper address_manager::translator by yourself
  *  So, just call address_translator_manager::singleton().translate(p) from your address_manager::translator and that's it.
  *  It'll translate addresses based on 'address_translator' objects, when one gets constructed it turns into a possible translator.
  *  At the constructor of your derived 'address_translator' make the map object to have [addr_to_translate] = translated_addr;
  *  There's also the virtual method 'fallback' that will get called when the translation wasn't possible, you can do some fallback stuff here
  *      (such as return the pointer as is or output a error message)
  */

#include "pch.h"

namespace Injector
{
	/*
	 *  address_translator
	 *      Base for an address translator
	 */
	class address_translator
	{
	private:
		bool enabled;
		void add();
		void remove();

	protected:
		friend class address_translator_manager;
		std::map<memory_pointer_raw, memory_pointer_raw> map;

	public:
		address_translator() : enabled(true)
		{
			// Must have bounds filled with min ptr and max ptr to have search working properly
			map.insert(std::make_pair(raw_ptr(0x00000000u), raw_ptr(0x00000000u)));
			map.insert(std::make_pair(raw_ptr(0xffffffffu), raw_ptr(0xffffffffu)));
			add();
		}

		~address_translator()
		{
			remove();
		}

		virtual void* fallback(void*) const
		{
			return nullptr;
		}


		// Enables or disables this translator
		void enable(bool enable_it)
		{
			if (enable_it) this->enable();
			else          this->disable();
		}

		// Enables this translator
		void enable()
		{
			this->enabled = true;
		}

		// Disables this translator
		void disable()
		{
			this->enabled = false;
		}

		// Checks if this translator is enabled
		bool is_enabled() const
		{
			return enabled;
		}
	};

	/*
	 *  address_translator_manager
	 *      Manages the address_translator objects
	 */
	class address_translator_manager
	{
	protected:
		friend class address_manager;
		friend class address_translator;

		std::list<const address_translator*> translators;

		void add(const address_translator& t)
		{
			translators.push_front(&t);
		}

		void remove(const address_translator& t)
		{
			translators.remove(&t);
		}

	public:
		// Translates the address p
		void* translator(void* p);

		// Singleton object
		static address_translator_manager& singleton()
		{
			static address_translator_manager mgr;
			return mgr;
		}
	};



	inline void* address_translator_manager::translator(void* p_)
	{
		static const size_t max_ptr_dist = 7;

		// Tries to find an address in a translator map
		auto try_map = [](const std::map<memory_pointer_raw, memory_pointer_raw>& map, memory_pointer_raw p) -> memory_pointer_raw
		{
			memory_pointer_raw result = nullptr;

			// Find first element in the map that is greater than or equal to p
			auto it = map.lower_bound(p);
			if (it != map.end())
			{
				// If it's not exactly the address, get back one position on the table
				if (it->first != p) --it;

				auto diff = (p - it->first).as_int();       // What's the difference between p and that address?
				if (diff <= max_ptr_dist)                    // Could we live with this difference in hands?
					result = it->second + raw_ptr(diff);    // Yes, we can!
			}

			return result;
		};


		//
		memory_pointer_raw result = nullptr;

		// Try to find translation for this pointer
		auto& mgr = address_translator_manager::singleton().translators;
		for (auto it = mgr.begin(); result == nullptr && it != mgr.end(); ++it)
		{
			auto& t = **it;
			if (t.is_enabled()) result = try_map(t.map, p_);
		}

		// If we couldn't translate the address, notify and try to fallback
		if (result.is_null())
		{
			for (auto it = mgr.begin(); result == nullptr && it != mgr.end(); ++it)
			{
				auto& t = **it;
				if (t.is_enabled())  result = t.fallback(p_);
			}
		}

		return result.get();
	}

	inline void address_translator::add()
	{
		address_translator_manager::singleton().add(*this);
	}

	inline void address_translator::remove()
	{
		address_translator_manager::singleton().remove(*this);
	}
}
